# Import python packages
import streamlit as st

# ======= BEGIN VANNA SETUP =======
from vanna.ZhipuAI.ZhipuAI_Chat import ZhipuAI_Chat
from vanna.chromadb.chromadb_vector import ChromaDB_VectorStore


class ZhipuAIVanna(ChromaDB_VectorStore, ZhipuAI_Chat):
    def __init__(self, config=None):
        ChromaDB_VectorStore.__init__(self, config=config)
        ZhipuAI_Chat.__init__(self, config=config)


vn = ZhipuAIVanna(

    config={"model": "glm-4", "api_key": "d3708ee404327e207b2f003775e06908.X3dgRCxbkyDfEIbh"})  # chatglm2-6b

vn.connect_to_mysql(host='10.240.53.162',
                    dbname='test',
                    user='rdsroot',
                    password='Geely@db20211206',
                    port=3306)
vn.train(ddl="""
        CREATE TABLE IF NOT EXISTS users (
            id INT PRIMARY KEY COMMENT '用户ID'， 
            username VARCHAR(50) COMMENT '用户名'， 
            email VARCHAR(100) COMMENT '电子邮件'， 
            age INT COMMENT '年龄'， 
            gender VARCHAR(10) COMMENT '性别（男/女）'， 
            city VARCHAR(50) COMMENT '城市'
        ) COMMENT='用户信息表' CHARACTER SET=utf8mb4 COLLATE=utf8mb4_unicode_ci;
        """)

# ======= END VANNA SETUP =======

my_question = st.session_state.get("my_question", default=None)

if my_question is None:
    my_question = st.text_input(
        "Ask me a question about your data",
        key="my_question",
    )
else:
    st.text(my_question)

    sql = vn.generate_sql(my_question)

    st.text(sql)

    df = vn.run_sql(sql)

    st.dataframe(df, use_container_width=True)

    code = vn.generate_plotly_code(question=my_question, sql=sql, df=df)

    fig = vn.get_plotly_figure(plotly_code=code, df=df)

    st.plotly_chart(fig, use_container_width=True)
